from datetime import time

from vnpy_ctastrategy import (
    CtaTemplate,
    StopOrder,
    TickData,
    BarData,
    TradeData,
    OrderData,
    ArrayManager,
)
from vnpy.trader.constant import Direction

from daily_generator import DailyBarGenerator
from result import ResultManager, TradeResult


class TurtleLimitSignal(CtaTemplate):
    """海龟信号"""

    author: str = "用Python的交易员"

    # 参数
    short_entry_window: int = 18
    short_exit_window: int = 7
    long_entry_window: int = 55
    long_exit_window: int = 25
    n_window: int = 20
    trail_window: int = 20
    trading_size: int = 1
    contract_size: int = 100
    price_add: int = 10

    # 变量
    target: int = 0             # 目标仓位

    # 名称列表
    parameters = [
        "short_entry_window",
        "short_exit_window",
        "long_entry_window",
        "long_exit_window",
        "n_window",
        "trail_window",
        "trading_size",
        "contract_size",
        "price_add"
    ]
    variables = [
        "target"
    ]

    def __init__(self, cta_engine, strategy_name, vt_symbol, setting) -> None:
        """构造函数"""
        super().__init__(cta_engine, strategy_name, vt_symbol, setting)

        self.short_factor: TurtleFactor = TurtleFactor(
            entry_window=self.short_entry_window,
            exit_window=self.short_exit_window,
            n_window=self.n_window,
            trail_window=self.trail_window,
            trading_size=self.trading_size,
            contract_size=self.contract_size,
            unit_limit=1,
            pnl_filter=True
        )

        self.long_factor: TurtleFactor = TurtleFactor(
            entry_window=self.long_entry_window,
            exit_window=self.long_exit_window,
            n_window=self.n_window,
            trail_window=self.trail_window,
            trading_size=self.trading_size,
            contract_size=self.contract_size,
            unit_limit=1,
            pnl_filter=False
        )

        self.current_factor: TurtleFactor = None

    def on_init(self) -> None:
        """初始化"""
        self.write_log("策略初始化")
        self.load_bar(20)

    def on_start(self) -> None:
        """启动"""
        pass

    def on_stop(self) -> None:
        """停止"""
        pass

    def on_tick(self, tick: TickData) -> None:
        """Tick推送"""
        pass

    def on_bar(self, bar: BarData) -> None:
        """原始K线推送"""
        # 全撤之前委托
        self.cancel_all()

        # 推送给因子计算
        self.short_factor.on_bar(bar)
        self.long_factor.on_bar(bar)

        # 如果尚未开仓
        if not self.current_factor:
            short_target = self.short_factor.get_target()
            long_target = self.long_factor.get_target()

            # 优先考虑短周期因子，开仓后记录当前因子
            if short_target:
                self.target = short_target
                self.current_factor = self.short_factor
            elif long_target:
                self.target = long_target
                self.current_factor = self.long_factor
        else:
            self.target = self.current_factor.get_target()

            # 平仓后清空当前因子
            if not self.target:
                self.current_factor = None

        # 计算目标和实际仓位差
        diff: int = self.target - self.pos      

        # 基于仓位差执行交易
        if diff > 0:
            price: float = bar.close_price + self.price_add

            # 由于海龟所有开平仓都会先回到仓位0的情况
            # 因此只需要考虑本次是开仓还是平仓即可
            if self.pos < 0:
                self.cover(price, abs(diff))
            else:
                self.buy(price, abs(diff))
        elif diff < 0:
            price: float = bar.close_price - self.price_add

            if self.pos > 0:
                self.sell(price, abs(diff))
            else:
                self.short(price, abs(diff))

    def on_trade(self, trade: TradeData) -> None:
        """成交推送"""
        pass

    def on_order(self, order: OrderData) -> None:
        """委托推送"""
        pass

    def on_stop_order(self, stop_order: StopOrder) -> None:
        """停止单推送"""
        pass


class TurtleFactor:
    """海龟因子"""

    def __init__(
        self,
        entry_window: int,
        exit_window: int,
        n_window: int,
        trail_window: int,
        trading_size: int,
        contract_size: int,
        unit_limit: int,
        pnl_filter: bool
    ) -> None:
        """构造函数"""
        # 参数
        self.entry_window: int = entry_window
        self.exit_window: int = exit_window
        self.n_window: int = n_window
        self.trail_window: int = trail_window
        self.trading_size: int = trading_size
        self.contract_size: int = contract_size
        self.unit_limit: int = unit_limit
        self.pnl_filter: bool = pnl_filter

        # 变量
        self.entry_up: float = 0.0       # 入场通道
        self.entry_down: float = 0.0

        self.exit_up: float = 0.0        # 出场通道
        self.exit_down: float = 0.0

        self.n: float = 0.0              # 波动度量

        # 移动止损参考：trail_window周期最高/最低
        self.trail_up: float = 0.0
        self.trail_down: float = 0.0

        self.long_entry: float = 0.0     # 开仓价格
        self.short_entry: float = 0.0

        self.target: int = 0             # 目标仓位
        self.traded: bool = False        # 日内交易过

        # 工具
        self.am = ArrayManager()
        self.bg = DailyBarGenerator(self.on_daily_bar, time(14, 59))
        self.rm = ResultManager(self.contract_size)

    def on_bar(self, bar: BarData):
        """原始K线推送"""
        # 每日只允许交易一次
        if self.am.inited and not self.traded:
            old_target = self.target
        
            # 判断当前目标
            if not self.target:
                self.check_long_target(bar)
                self.check_short_target(bar)
            elif self.target > 0:
                self.check_long_target(bar)

                # 计算固定止损价格
                long_stop: float = self.long_entry - 2 * self.n
                # 和离场通道与移动止损比较，取更高价挂出停止单
                trail_long_stop: float = self.trail_up - 2 * self.n
                long_stop = max(long_stop, self.exit_down)
                long_stop = max(long_stop, trail_long_stop)

                if bar.low_price <= long_stop:
                    self.target = 0
            elif self.target < 0:
                self.check_short_target(bar)

                short_stop: float = self.short_entry + 2 * self.n
                # 和离场通道与移动止损比较，取更低价挂出停止单
                trail_short_stop: float = self.trail_down + 2 * self.n
                short_stop = min(short_stop, self.exit_up)
                short_stop = min(short_stop, trail_short_stop)

                if bar.high_price >= short_stop:
                    self.target = 0

            # 记录今天已经执行过交易
            if old_target != self.target:
                self.traded = True

                # 创建成交对象并记录
                diff = self.target - old_target

                if diff > 0:
                    direction = Direction.LONG
                else:
                    direction = Direction.SHORT

                trade = TradeData(
                    gateway_name="FACTOR",
                    symbol=bar.symbol,
                    exchange=bar.exchange,
                    orderid=str(bar.datetime),
                    tradeid=str(bar.datetime),
                    direction=direction,
                    price=bar.close_price,
                    volume=abs(diff),
                    datetime=bar.datetime
                )
                self.rm.update_trade(trade)

        self.bg.update_bar(bar)

    def on_daily_bar(self, bar: BarData):
        """日K线推送"""
        # 缓存K线序列
        self.am.update_bar(bar)
        if not self.am.inited:
            return        

        # 只有无持仓时，才更新入场通道位置和波动度量
        if not self.target:
            self.entry_up, self.entry_down = self.am.donchian(self.entry_window)
            self.n = self.am.atr(self.n_window)

        self.exit_up, self.exit_down = self.am.donchian(self.exit_window)

        # 计算trail_window窗口的最高/最低用于移动止损
        self.trail_up, self.trail_down = self.am.donchian(self.trail_window)

        # 新的一天清空交易记录
        self.traded = False

    def check_long_target(self, bar: BarData):
        """检查多头"""
        level: int = self.unit_limit

        while level > 0:
            level_limit: int = level * self.trading_size
            level_price = self.entry_up + self.n * (level - 1) * 0.5

            if self.target < level_limit and bar.high_price >= level_price:
                self.target = level_limit
                self.long_entry = level_price

            level -= 1

    def check_short_target(self, bar: BarData):
        """检查空头"""
        level: int = self.unit_limit

        while level > 0:
            level_limit: int = -level * self.trading_size
            level_price = self.entry_down - self.n * (level - 1) * 0.5

            if self.target > level_limit and bar.low_price <= level_price:
                self.target = level_limit
                self.short_entry = level_price

            level -= 1

    def get_target(self) -> int:
        """获取因子目标"""
        # 如果不进行过滤，则直接返回目标
        if not self.pnl_filter:
            return self.target

        # 查询逐笔对冲盈亏记录
        results: list[TradeResult] = self.rm.get_results()

        # 检查上一笔交易如果盈利，则本次信号忽略
        if results:
            last_result: TradeResult = results[-1]
            if last_result.pnl > 0:
                return 0

        return self.target
